#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <memory.h>

#define PI 3.14159265358979323846

void swap(double *a, double *b)
{
    double t = *a;
    *a = *b;
    *b = t;
}

void bitrev(double *real, double *imag, int n)
{
    int j = 0;
    int m = 0;
    for (int i = 0; i < n; i++)
    {
        if (j > i)
        {
            swap(&real[i], &real[j]);
            swap(&imag[i], &imag[j]);
        }
        m = n / 2;
        while (j >= m && m != 0)
        {
            j -= m;
            m >>= 1;
        }
        j += m;
    }
}

int log2int(int n)
{
    int m = 0;
    for (int i = n; i > 1; i = i >> 1, m++);
    return m;
}

void fft(double *real, double *imag, int n)
{
    int M = log2int(n);

    bitrev(real, imag, n);

    for (int m = 0; m < M; m++)
    {
        int K = n >> (m + 1);
        for (int k = 0; k < K; k++)
        {
            int J = 2 << m;
            int base = k * J;
            for (int j = 0; j < J / 2; j++)
            {
                int t = base + j;
                double wnr = real[t + J / 2] * cos(-2 * PI * j * K / n) - imag[t + J / 2] * sin(-2 * PI * j * K / n);
                double wni = real[t + J / 2] * sin(-2 * PI * j * K / n) + imag[t + J / 2] * cos(-2 * PI * j * K / n);
                double tempr = real[t];
                double tempi = imag[t];
                real[t] += wnr;
                imag[t] += wni;
                real[t + J / 2] = tempr - wnr;
                imag[t + J / 2] = tempi - wni;
            }
        }
    }
}

void ifft(double *real, double *imag, int n)
{
    int M = log2int(n);

    bitrev(real, imag, n);

    for (int m = 0; m < M; m++)
    {
        int K = n >> (m + 1);
        for (int k = 0; k < K; k++)
        {
            int J = 2 << m;
            int base = k * J;
            for (int j = 0; j < J / 2; j++)
            {
                int t = base + j;
                double wnr = real[t + J / 2] * cos(2 * PI * j * K / n) - imag[t + J / 2] * sin(2 * PI * j * K / n);
                double wni = real[t + J / 2] * sin(2 * PI * j * K / n) + imag[t + J / 2] * cos(2 * PI * j * K / n);
                double tempr = real[t];
                double tempi = imag[t];
                real[t] += wnr;
                imag[t] += wni;
                real[t + J / 2] = tempr - wnr;
                imag[t + J / 2] = tempi - wni;
            }
        }
    }

    for (int i = 0; i < n; i++)
    {
        real[i] /= n;
        imag[i] /= n;
    }
}

void dft(double *real, double *imag, int n)
{
    double *out_real = (double *)malloc(n * sizeof(double));
    double *out_imag = (double *)malloc(n * sizeof(double));

    memset(out_real, 0, n * sizeof(double));
    memset(out_imag, 0, n * sizeof(double));

    for (int k = 0; k < n; k++)
    {
        for (int m = 0; m < n; m++)
        {
            out_real[k] += (real[m] * cos(-2 * PI * m * k / n) - imag[m] * sin(-2 * PI * m *  k / n));
            out_imag[k] += (real[m] * sin(-2 * PI * m * k / n) + imag[m] * cos(-2 * PI * m *  k / n));
        }
    }

    memcpy(real, out_real, n * sizeof(double));
    memcpy(imag, out_imag, n * sizeof(double));

    free(out_real);
    free(out_imag);
}

void idft(double *real, double *imag, int n)
{
    double *out_real = (double *)malloc(n * sizeof(double));
    double *out_imag = (double *)malloc(n * sizeof(double));

    memset(out_real, 0, n * sizeof(double));
    memset(out_imag, 0, n * sizeof(double));

    for (int k = 0; k < n; k++)
    {
        for (int m = 0; m < n; m++)
        {
            out_real[k] += (real[m] * cos(2 * PI * m * k / n) - imag[m] * sin(2 * PI * m *  k / n));
            out_imag[k] += (real[m] * sin(2 * PI * m * k / n) + imag[m] * cos(2 * PI * m *  k / n));
        }
    }

    for (int i = 0; i < n; i++)
    {
        out_real[i] /= n;
        out_imag[i] /= n;
    }

    memcpy(real, out_real, n * sizeof(double));
    memcpy(imag, out_imag, n * sizeof(double));

    free(out_real);
    free(out_imag);
}

void show(double *real, double *imag, int n)
{
    for (int i = 0; i < n; i++) 
    {
        printf("%lf%+lfj\n", real[i], imag[i]);
    }
}


// test
#include <time.h>
#define N (1024 * 8)

int main(void)
{    
    int n = 128;
    double real[N] = {0};
    double imag[N] = {0};

    // test dft
    if (0)
    { 
        for (int i = 0; i < n; i++) 
        {       
            real[i] = (double)i;
            imag[i] = (double)0;
        }    
        clock_t s = clock();
        dft(real, imag, n);
        idft(real, imag, n);
        clock_t e = clock();
        show(real, imag, n);
        printf("run %f seconds\n", (double)(e - s) / CLOCKS_PER_SEC);
    }
    
    // test fft
    if (1)
    {
        for (int i = 0; i < n; i++) 
        {       
            real[i] = (double)i;
            imag[i] = (double)0;
        }    
        clock_t s = clock();
        fft(real, imag, n);
        ifft(real, imag, n);
        clock_t e = clock();
        show(real, imag, n);
        printf("run %f seconds\n", (double)(e - s) / CLOCKS_PER_SEC);
    }

    return 0;
}
